"""
FinRL框架主程序文件
提供训练、测试和交易三种模式的统一入口
"""

from __future__ import annotations

import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

from argparse import ArgumentParser
from typing import List

# 从配置文件导入各种参数和设置
from finrl.config import ALPACA_API_BASE_URL
from finrl.config import DATA_SAVE_DIR
from finrl.config import ERL_PARAMS
from finrl.config import INDICATORS
from finrl.config import LIHUA_INDICATORS
from finrl.config import RESULTS_DIR
from finrl.config import TENSORBOARD_LOG_DIR
from finrl.config import TEST_END_DATE
from finrl.config import TEST_START_DATE
from finrl.config import TRADE_END_DATE
from finrl.config import TRADE_START_DATE
from finrl.config import TRAIN_END_DATE
from finrl.config import TRAIN_START_DATE
from finrl.config import TRAINED_MODEL_DIR
from finrl.config_tickers import DOW_30_TICKER
from finrl.config_tickers import CSI_300_TICKER

from finrl.config_tickers import SSE_50_TICKER
from finrl.meta.env_stock_trading.env_stocktrading_np import StockTradingEnv

# construct environment

# try:
#     from finrl.config_private import ALPACA_API_KEY, ALPACA_API_SECRET
# except ImportError:
#     raise FileNotFoundError(
#         "Please set your own ALPACA_API_KEY and ALPACA_API_SECRET in config_private.py"
#     )


def build_parser():
    """
    构建命令行参数解析器
    Returns:
        ArgumentParser: 配置好的参数解析器
    """
    parser = ArgumentParser()
    parser.add_argument(
        "--mode",
        dest="mode",
        help="启动模式: train(训练), test(测试), trade(交易)",
        metavar="MODE",
        default="train",
    )
    return parser


def check_and_make_directories(directories: list[str]):
    """
    检查并创建必要的目录
    
    Args:
        directories (list[str]): 需要检查和创建的目录列表
    """
    for directory in directories:
        if not os.path.exists(directory):
            os.makedirs(directory)


def main() -> int:
    """
    主函数，根据命令行参数执行不同的操作模式
    
    Returns:
        int: 程序退出码，0表示正常退出
    """
    # 解析命令行参数
    parser = build_parser()
    options = parser.parse_args()
    
    # 检查并创建必要的数据存储目录
    check_and_make_directories(
        [DATA_SAVE_DIR, TRAINED_MODEL_DIR, TENSORBOARD_LOG_DIR, RESULTS_DIR]
    )

    # 根据模式执行不同操作
    if options.mode == "train":
        # 训练模式
        from finrl import train

        env = StockTradingEnv

        # elegantrl的示例配置
        # 当前在meta中，对于yahoofinance，kwargs是{}。对于其他数据源，如joinquant，kwargs不为空
        kwargs = (
            {}
        )  
        train(
            start_date=TRAIN_START_DATE,        # 训练开始日期
            end_date=TRAIN_END_DATE,            # 训练结束日期
            #ticker_list=DOW_30_TICKER,         # 股票列表(DOW 30)
            #data_source="yahoofinance",                # 数据源
            #technical_indicator_list=INDICATORS,  # 技术指标列表
            ticker_list=SSE_50_TICKER,
            technical_indicator_list=LIHUA_INDICATORS,  # 技术指标列表
            data_source="lihua",
            time_interval="1D",                 # 时间间隔(1天) 日线

            drl_lib="elegantrl",                # 深度强化学习库
            env=env,                            # 环境类
            model_name="ppo",                   # 模型名称(PPO)
            cwd="./test_ppo",                   # 当前工作目录
            erl_params=ERL_PARAMS,              # ERL参数
            break_step=5e5,                     # 最大训练步数
            kwargs=kwargs,                      # 其他参数
        )

        # 测试模式
        from finrl import test

        envTest = StockTradingEnv

        # elegantrl的示例配置
        # 当前在meta中，对于yahoofinance，kwargs是{}。对于其他数据源，如joinquant，kwargs不为空
        kwargs = {}

        account_value_erl = test(  # noqa
            start_date=TEST_START_DATE,  # 测试开始日期
            end_date=TEST_END_DATE,  # 测试结束日期
            # ticker_list=DOW_30_TICKER,          # 股票列表(DOW 30)
            # data_source="yahoofinance",         # 数据源
            # technical_indicator_list=INDICATORS,  # 技术指标列表
            ticker_list=SSE_50_TICKER,
            technical_indicator_list=LIHUA_INDICATORS,  # 技术指标列表
            data_source="lihua",
            time_interval="1D",  # 时间间隔(1天)

            drl_lib="elegantrl",  # 深度强化学习库
            env=envTest,  # 环境类
            model_name="ppo",  # 模型名称(PPO)
            cwd="./test_ppo",  # 当前工作目录
            net_dimension=512,  # 网络维度
            kwargs=kwargs,  # 其他参数
        )

    elif options.mode == "test":
        # 测试模式
        from finrl import test

        env = StockTradingEnv

        # elegantrl的示例配置
        # 当前在meta中，对于yahoofinance，kwargs是{}。对于其他数据源，如joinquant，kwargs不为空
        kwargs = {}

        account_value_erl = test(  # noqa
            start_date=TEST_START_DATE,         # 测试开始日期
            end_date=TEST_END_DATE,             # 测试结束日期
            # ticker_list=DOW_30_TICKER,          # 股票列表(DOW 30)
            # data_source="yahoofinance",         # 数据源
            # technical_indicator_list=INDICATORS,  # 技术指标列表
            ticker_list=SSE_50_TICKER,
            technical_indicator_list=LIHUA_INDICATORS,  # 技术指标列表
            data_source="lihua",
            time_interval="1D",                 # 时间间隔(1天)

            drl_lib="elegantrl",                # 深度强化学习库
            env=env,                            # 环境类
            model_name="ppo",                   # 模型名称(PPO)
            cwd="./test_ppo1",                   # 当前工作目录
            net_dimension=512,                  # 网络维度
            kwargs=kwargs,                      # 其他参数
        )
    elif options.mode == "trade":
        # 交易模式
        from finrl import trade

        # 导入Alpaca API密钥，如果不存在则抛出异常
        try:
            from finrl.config_private import ALPACA_API_KEY, ALPACA_API_SECRET
        except ImportError:
            raise FileNotFoundError(
                "Please set your own ALPACA_API_KEY and ALPACA_API_SECRET in config_private.py"
            )
        env = StockTradingEnv
        kwargs = {}
        trade(
            start_date=TRADE_START_DATE,        # 交易开始日期
            end_date=TRADE_END_DATE,            # 交易结束日期
            ticker_list=DOW_30_TICKER,          # 股票列表(DOW 30)
            data_source="yahoofinance",         # 数据源
            time_interval="1D",                 # 时间间隔(1天)
            technical_indicator_list=INDICATORS, # 技术指标列表
            drl_lib="elegantrl",                # 深度强化学习库
            env=env,                            # 环境类
            model_name="ppo",                   # 模型名称(PPO)
            API_KEY=ALPACA_API_KEY,             # Alpaca API密钥
            API_SECRET=ALPACA_API_SECRET,       # Alpaca API密钥
            API_BASE_URL=ALPACA_API_BASE_URL,   # Alpaca API基础URL
            trade_mode="paper_trading",         # 交易模式(纸面交易)
            if_vix=True,                        # 是否使用VIX指数
            kwargs=kwargs,                      # 其他参数
            state_dim=len(DOW_30_TICKER) * (len(INDICATORS) + 3)
            + 3,  # bug fix: for ppo add dimension of state/observations space =  len(stocks)* len(INDICATORS) + 3+ 3*len(stocks)
            action_dim=len(
                DOW_30_TICKER
            ),  # bug fix: for ppo add dimension of action space = len(stocks)
        )
    else:
        # 无效模式
        raise ValueError("Wrong mode.")
    return 0


# 用户可以在终端输入以下命令:
# python main.py --mode=train    # 训练模式
# python main.py --mode=test     # 测试模式
# python main.py --mode=trade    # 交易模式
if __name__ == "__main__":
    raise SystemExit(main())